# -*- coding: utf-8 -*-
from abc import ABCMeta, abstractmethod
import os
import numpy as np
import trimesh
from .mean_curvature import discrete_mean_curvature_measure_on_mesh, sphere_ball_intersection
import matplotlib.pyplot as plt
from autolab_core import SimilarityTransform


class GraspableObject:
    """ Encapsulates geometric structures for computing contact in grasping.
    
    Attributes
    ----------
    sdf : :obj:`Sdf3D`
        signed distance field, for quickly computing contact points
    mesh : :obj:`Mesh3D`
        3D triangular mesh to specify object geometry, should match SDF
    key : :obj:`str`
        object identifier, usually given from the database
    model_name : :obj:`str`
        name of the object mesh as a .obj file, for use in collision checking
    mass : float
        mass of the object
    convex_pieces : :obj:`list` of :obj:`Mesh3D`
        convex decomposition of the object geom for collision checking
    """
    __metaclass__ = ABCMeta

    def __init__(self, mesh, key='', model_name='', mass=0.1):
        self.mesh_ = mesh
        self.key_ = key
        self.mesh_tree = self.mesh_.kdtree
        self.model_name_ = model_name # for OpenRave usage, gross!
        self.mass_ = mesh.mass if mesh.mass is not None else mass

    @property
    def mesh(self):
        return self.mesh_

    @property
    def mass(self):
        return self.mass_

    @property
    def key(self):
        return self.key_

    @property
    def model_name(self):
        return self.model_name_

    # @property
    # def convex_pieces(self):
    #     return self.convex_pieces_


class GraspableObject3D(GraspableObject):
    """ 3D Graspable object for computing contact in grasping.
    
    Attributes
    ----------
    sdf : :obj:`Sdf3D`
        signed distance field, for quickly computing contact points
    mesh : :obj:`Mesh3D`
        3D triangular mesh to specify object geometry, should match SDF
    key : :obj:`str`
        object identifier, usually given from the database
    model_name : :obj:`str`
        name of the object mesh as a .obj file, for use in collision checking
    mass : float
        mass of the object
    convex_pieces : :obj:`list` of :obj:`Mesh3D`
        convex decomposition of the object geom for collision checking
    """
    def __init__(self, mesh, key='',
                 model_name='', mass=0.1,
                 file_path=None,
                 threshold=80,
                 radius=0.003):
        if not isinstance(mesh, trimesh.Trimesh):
            raise ValueError('Must initialize 3D graspable object with Trimesh')

        GraspableObject.__init__(self, mesh, key=key,
                                 model_name=model_name, 
                                 mass=mass)
        assert file_path is not None
        self.mean_curvature_radius = radius
        self.file_path = file_path
        self.threshold = threshold
        self.mean_curvature = self.get_mean_curvature()
        self.badpoints_mask = self.get_badpoints_mask()

    def get_mean_curvature(self):
        curvature_file_path = self.file_path.replace('{}.stl'.format(self.key), '{}.npy'.format(self.key))
        if not os.path.exists(curvature_file_path):
            MeanCurvature = np.array(discrete_mean_curvature_measure_on_mesh(self.mesh, self.mean_curvature_radius) /
                                     sphere_ball_intersection(1, self.mean_curvature_radius))
        else:
            MeanCurvature = np.load(curvature_file_path)
        return MeanCurvature

    def get_badpoints_mask(self):
        mask_1 = -self.threshold < self.mean_curvature
        mask_2 = self.mean_curvature < self.threshold
        mask = mask_1 & mask_2
        return mask

    @property
    def diagonal_length(self):
        extents = self.mesh.bounding_box.primitive.extents
        
        return np.linalg.norm(extents) / 2

    def moment_arm(self, x):
        """ Computes the moment arm to a point x.

        Parameters
        ----------
        x : 3x1 :obj:`numpy.ndarray`
            point to get moment arm for
        
        Returns
        -------
        3x1 :obj:`numpy.ndarray`
        """
        return x - self.mesh.center_mass

    def rescale(self, scale):
        """ Rescales uniformly by a given factor.

        Parameters
        ----------
        scale : float
            the amount to scale the object

        Returns
        -------
        :obj:`GraspableObject3D`
            the graspable object rescaled by the given factor
        """
        stf = SimilarityTransform(scale=scale)
        mesh_rescaled = self.mesh_.transform(stf)
        return GraspableObject3D(mesh_rescaled, key=self.key,
                                 model_name=self.model_name, mass=self.mass)

    def transform(self, delta_T):
        """ Transform by a delta transform.


        Parameters
        ----------
        delta_T : :obj:`RigidTransform`
            the transformation from the current reference frame to the alternate reference frame
        
        Returns
        -------
        :obj:`GraspableObject3D`
             graspable object trasnformed by the delta
        """
        mesh_tf = self.mesh_.transform(delta_T)
        return GraspableObject3D(mesh_tf, key=self.key,
                                 model_name=self.model_name, mass=self.mass)

    def surface_information(self, grasp, width, num_steps, plot=False, direction1=None, direction2=None):
        """ Returns the patches on this object for a given grasp.

        Parameters
        ----------
        grasp : :obj:`ParallelJawPtGrasp3D`
            grasp to get the patch information for
        width : float
            width of jaw opening
        num_steps : int
            number of steps
        plot : bool
            whether to plot the intermediate computation, for debugging
        direction1 : normalized 3x1 :obj:`numpy.ndarray`
            direction along which to compute the surface information for the first jaw, if None then defaults to grasp axis
        direction2 : normalized 3x1 :obj:`numpy.ndarray`
            direction along which to compute the surface information for the second jaw, if None then defaults to grasp axis
       
        Returns
        -------
        :obj:`list` of :obj:`SurfaceWindow`
             surface patches, one for each contact
       """
        contacts_found, contacts = grasp.close_fingers(self)
        if not contacts_found:
            raise ValueError('Failed to find contacts')
        contact1, contact2 = contacts

        if plot:
            plt.figure()
            contact1.plot_friction_cone()
            contact2.plot_friction_cone()

            ax = plt.gca(projection='3d')

        window1 = contact1.surface_information(width, num_steps, direction=direction1)
        window2 = contact2.surface_information(width, num_steps, direction=direction2)
        return window1, window2, contact1, contact2

